# -*- coding: utf-8 -*-
"""
文件: tmdb_model_training.py
功能: 使用 5 种模型对 TMDB 电影数据进行评分预测
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.svm import SVR
from sklearn.metrics import mean_squared_error, r2_score
import xgboost as xgb
import warnings
import matplotlib.pyplot as plt
import seaborn as sns

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# ========================
# 1. 加载清洗后的数据
# ========================
try:
    df = pd.read_csv('tmdb_credits_cleaned.csv')
except FileNotFoundError:
    raise FileNotFoundError("请确保 'tmdb_credits_cleaned.csv' 文件在当前目录下！")

print(f"[OK] 数据形状: {df.shape}")
print(f"[INFO] 列名: {list(df.columns)}\n")

# ========================
# 2. 特征工程
# ========================
# 提取数值型特征
features = [
    'cast_size', 'crew_size', 'total_team_size',
    'num_top_actors', 'num_writers', 'has_female_actor'
]

# 创建新特征：是否为知名导演（如 James Cameron）
df['is_famous_director'] = df['director'].isin(['James Cameron', 'Christopher Nolan', 'Quentin Tarantino']).astype(int)

# 添加更多特征
df['team_ratio'] = df['cast_size'] / (df['crew_size'] + 1)  # 防止除零
df['writer_count'] = df['writers'].apply(len)


# 安全计算演员性别多样性
def calculate_gender_diversity(cast_list):
    if not isinstance(cast_list, list):
        return 0.0
    if len(cast_list) == 0:
        return 0.0
    female_count = sum(1 for actor in cast_list if isinstance(actor, dict) and actor.get('gender') == 1)
    return female_count / len(cast_list)


df['actor_gender_diversity'] = df['cast'].apply(calculate_gender_diversity)

# 确保所有特征列存在
required_cols = ['num_top_actors', 'num_writers', 'actor_gender_diversity']
for col in required_cols:
    if col not in df.columns:
        print(f"⚠️ 警告: 列 '{col}' 不存在，已创建默认值")
        df[col] = 0

# ========================
# 3. 目标变量处理
# ========================
# 检查并获取真实评分字段
target_col = None
possible_targets = ['vote_average', 'rating', 'score', 'avg_rating']

for col in possible_targets:
    if col in df.columns:
        target_col = col
        break

if target_col is None:
    # 如果没有评分字段，创建一个模拟评分字段用于演示
    print("⚠️ 未找到评分字段，创建模拟评分数据用于演示")
    np.random.seed(42)
    df['vote_average'] = np.random.uniform(1.0, 10.0, len(df))
    target_col = 'vote_average'

print(f"🎯 使用目标变量: {target_col}")

# ========================
# 4. 特征工程与选择
# ========================
# 选择最终特征
selected_features = features + ['is_famous_director', 'team_ratio', 'writer_count', 'actor_gender_diversity']

# 确保所有特征列都是数值类型
for col in selected_features:
    if col in df.columns:
        # 尝试转换为数值类型
        df[col] = pd.to_numeric(df[col], errors='coerce')

# 只选择存在的数值型特征
X = df.select_dtypes(include=[np.number])
# 移除目标变量列（如果存在）
if target_col in X.columns:
    X = X.drop(target_col, axis=1)

print(f"[INFO] 最终选择的特征列: {list(X.columns)}")
print(f"[INFO] 特征数据类型:\n{X.dtypes}\n")

# 设置目标变量
y = df[target_col]

# ========================
# 5. 数据预处理
# ========================
# 处理缺失值 - 更全面的缺失值处理
print("🔍 缺失值统计:")
missing_before = X.isnull().sum().sum()
print(f"  填充前NaN值总数: {missing_before}")

# 首先对数值型列使用平均值填充
X = X.fillna(X.mean(numeric_only=True))

# 对于仍然存在的NaN值（可能是在非数值列或全为NaN的列），使用中位数填充
X = X.fillna(X.median(numeric_only=True))

# 最后，如果还有NaN值，使用0填充
X = X.fillna(0)

# 再次检查是否还有NaN值
missing_after = X.isnull().sum().sum()
print(f"  填充后NaN值总数: {missing_after}")

if missing_after > 0:
    print("  ⚠️  警告: 仍有NaN值存在，将强制替换为0")
    X = X.fillna(0)

# 分割训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

print(f"📈 训练集大小: {X_train.shape}, 测试集大小: {X_test.shape}\n")

# ========================
# 5. 模型训练与评估
# ========================

models = {
    "线性回归": LinearRegression(),
    "随机森林": RandomForestRegressor(n_estimators=100, random_state=42),
    "XGBoost": xgb.XGBRegressor(n_estimators=100, random_state=42),
    "支持向量机": SVR(kernel='rbf'),
    "梯度提升": GradientBoostingRegressor(random_state=42)
}

results = {}
predictions = {}

for name, model in models.items():
    print(f"\n🚀 正在训练模型: {name}")
    try:
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)

        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)

        results[name] = {'MSE': mse, 'R²': r2}
        predictions[name] = y_pred
        print(f"  MSE: {mse:.4f}, R²: {r2:.4f}")
    except Exception as e:
        print(f"  ❌ 训练失败: {e}")

# ========================
# 6. 图表可视化
# ========================

# 6.1 模型性能对比柱状图
model_names = list(results.keys())
mse_values = [results[name]['MSE'] for name in model_names]
r2_values = [results[name]['R²'] for name in model_names]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# MSE对比
bars1 = ax1.bar(model_names, mse_values, color='skyblue')
ax1.set_title('各模型MSE对比')
ax1.set_ylabel('MSE (越小越好)')
ax1.tick_params(axis='x', rotation=45)

# 在柱状图上添加数值标签
for bar, value in zip(bars1, mse_values):
    ax1.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
             f'{value:.3f}', ha='center', va='bottom')

# R²对比
bars2 = ax2.bar(model_names, r2_values, color='lightcoral')
ax2.set_title('各模型R²对比')
ax2.set_ylabel('R² (越大越好)')
ax2.tick_params(axis='x', rotation=45)

# 在柱状图上添加数值标签
for bar, value in zip(bars2, r2_values):
    ax2.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
             f'{value:.3f}', ha='center', va='bottom')

plt.tight_layout()
plt.savefig('model_performance_comparison.png', dpi=300, bbox_inches='tight')
plt.show()

# 6.2 预测值vs真实值散点图（最佳模型）
best_model_name = min(results.keys(), key=lambda k: results[k]['MSE'])
best_predictions = predictions[best_model_name]

plt.figure(figsize=(10, 8))
plt.scatter(y_test, best_predictions, alpha=0.6)
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
plt.xlabel('真实值')
plt.ylabel('预测值')
plt.title(f'{best_model_name} - 预测值 vs 真实值')
plt.grid(True, alpha=0.3)
plt.savefig('prediction_vs_actual.png', dpi=300, bbox_inches='tight')
plt.show()

# 6.3 特征重要性图（针对树模型）
tree_models = ["随机森林", "XGBoost", "梯度提升"]
for model_name in tree_models:
    if model_name in models:
        model = models[model_name]
        if hasattr(model, 'feature_importances_'):
            plt.figure(figsize=(10, 6))
            importances = model.feature_importances_
            indices = np.argsort(importances)[::-1]
            feature_names = X.columns

            plt.bar(range(len(importances)), importances[indices])
            plt.title(f'{model_name} - 特征重要性')
            plt.xticks(range(len(importances)), [feature_names[i] for i in indices], rotation=45)
            plt.tight_layout()
            plt.savefig(f'{model_name}_feature_importance.png', dpi=300, bbox_inches='tight')
            plt.show()

# 6.4 目标变量分布直方图
plt.figure(figsize=(10, 6))
plt.hist(y, bins=30, edgecolor='black', alpha=0.7)
plt.xlabel(target_col)
plt.ylabel('频数')
plt.title(f'{target_col} 分布')
plt.grid(True, alpha=0.3)
plt.savefig('target_distribution.png', dpi=300, bbox_inches='tight')
plt.show()

# 6.5 特征相关性热力图
plt.figure(figsize=(12, 10))
correlation_matrix = pd.concat([X, y], axis=1).corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0,
            fmt='.2f', square=True)
plt.title('特征相关性热力图')
plt.tight_layout()
plt.savefig('feature_correlation_heatmap.png', dpi=300, bbox_inches='tight')
plt.show()

# ========================
# 7. 输出结果对比
# ========================
print("\n📋 所有模型性能对比:")
for name, metrics in results.items():
    print(f"{name}: MSE={metrics['MSE']:.4f}, R²={metrics['R²']:.4f}")

# ========================
# 8. 保存最佳模型
# ========================
print(f"\n🏆 最佳模型: {best_model_name}")

# 保存最佳模型（可选）
import joblib

joblib.dump(models[best_model_name], 'best_model.pkl')
print("💾 最佳模型已保存为 'best_model.pkl'")

# ========================
# 9. 生成模型报告
# ========================
print("\n📊 图表已生成并保存:")
print("1. model_performance_comparison.png - 模型性能对比图")
print("2. prediction_vs_actual.png - 最佳模型预测值vs真实值")
print("3. *_feature_importance.png - 树模型特征重要性图")
print("4. target_distribution.png - 目标变量分布图")
print("5. feature_correlation_heatmap.png - 特征相关性热力图")
